package top.atcra.rpc.discovery.register.impl;

import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.ZooKeeper;
import top.atcra.rpc.config.ConfigLoader;
import top.atcra.rpc.discovery.register.ServiceRegistry;
import top.atcra.rpc.exception.DiscoveryException;
import top.atcra.rpc.protection.health.ServiceUpAndDownWatcher;
import top.atcra.rpc.service.ServiceConfig;
import top.atcra.rpc.utils.IPAddrUtil;
import top.atcra.rpc.utils.ZookeeperUtil;

import java.net.InetSocketAddress;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;

/**
 * 基于Zookeeper的注册中心实现
 *
 * @author atcra
 * @date 2025-3-22
 */
@Slf4j
public class ZookeeperServiceRegistry implements ServiceRegistry {
    protected final ZooKeeper zk;


    public ZookeeperServiceRegistry(String connectAddr) {
        zk = ZookeeperUtil.createZookeeper(connectAddr);
    }

    /**
     * 注册服务，会在Zookeeper中维护如图所示的目录树
     * <br>
     * rpc-service：Base路径 <br>
     * ├ providers：服务提供方 <br>
     * │ └ service1 允许一个Host向多个网络端口发布服务 <br>
     * │	 └ dev	分组 <br>
     * │	   ├ 192.168.2.1:8090 <br>
     * │	   └ 192.168.193.1:8090 <br>
     * └ consumers <br>
     *
     * @param serviceDef 服务定义
     */
    @Override
    public void register(ServiceConfig serviceDef) {
        var cfg = ConfigLoader.getConfig();
        var zkcfg = cfg.getZookeeper();
        String serviceBasePath = zkcfg.getBase() + "/" + serviceDef.getName();
        String groupBasePath = serviceBasePath + "/" + serviceDef.getGroup();
        String localPath = groupBasePath + "/" + IPAddrUtil.getLocalIpv4InSubnet(cfg.getSubnetPrefix()) + ":" + cfg.getPort();

        ZookeeperUtil.createNodeIfNotExist(zk,zkcfg.getBase(),null,CreateMode.PERSISTENT);
        ZookeeperUtil.createNodeIfNotExist(zk, serviceBasePath, null, CreateMode.PERSISTENT);
        ZookeeperUtil.createNodeIfNotExist(zk, groupBasePath, null, CreateMode.PERSISTENT);
        ZookeeperUtil.createNodeIfNotExist(zk, localPath, null, CreateMode.EPHEMERAL);
        log.info("zookeeper register serviceDef: {}", serviceDef);
    }

    @Override
    public void register(ServiceConfig serviceDef, int duplicateCount) {
        var cfg = ConfigLoader.getConfig();
        var zkcfg = cfg.getZookeeper();
        String serviceBasePath = zkcfg.getBase() + "/" + serviceDef.getName();
        String groupBasePath = serviceBasePath + "/" + serviceDef.getGroup();
        String localPath = groupBasePath + "/" + IPAddrUtil.getLocalIpv4InSubnet(cfg.getSubnetPrefix()) + ":" + cfg.getPort();

        ZookeeperUtil.createNodeIfNotExist(zk, serviceBasePath, null, CreateMode.PERSISTENT);
        ZookeeperUtil.createNodeIfNotExist(zk, groupBasePath, null, CreateMode.PERSISTENT);
        for (int i = 0; i < duplicateCount; i++)
            ZookeeperUtil.createNodeIfNotExist(zk, localPath, null, CreateMode.EPHEMERAL);
        log.info("zookeeper register serviceDef: {},duplicate {} times.", serviceDef, duplicateCount);
    }

    Function<Void, ServiceUpAndDownWatcher> createUpAndDownWatcher;

    public ZookeeperServiceRegistry setCreateUpAndDownWatcher(Function<Void, ServiceUpAndDownWatcher> createUpAndDownWatcher) {
        this.createUpAndDownWatcher = createUpAndDownWatcher;
        return this;
    }

    protected static final Map<List<Object>, ServiceUpAndDownWatcher> watchers = new ConcurrentHashMap<>();

    @Override
    public List<InetSocketAddress> lookup(String serviceName, String group) {
        ServiceUpAndDownWatcher watcher = null;
        if (createUpAndDownWatcher != null) {
            // 避免为同一个Service重复创建watcher
            if (!watchers.containsKey(List.of(serviceName, group))) {
                watcher = createUpAndDownWatcher.apply(null);
                watchers.put(List.of(serviceName, group), watcher);
            }
        }

        try {
            var zkcfg = ConfigLoader.getConfig().getZookeeper();
            String groupPath = zkcfg.getBase() + "/" + serviceName + "/" + group;
            var list = zk.getChildren(groupPath, watcher).stream().map(
                    (addr) -> {
                        String[] split = addr.split(":");
                        return new InetSocketAddress(split[0], Integer.parseInt(split[1]));
                    }
            ).toList();
            if (list.isEmpty()) {
                throw new DiscoveryException(String.format("No %s(%s) can be found by zookeeper lookup service: ", serviceName, group));
            } else
                log.info("zookeeper lookup service by name({}) and group({}). Result is {}", serviceName, group, list);
            return list;
        } catch (Exception e) {
            log.error("zookeeper lookup serviceDef failed: {}", serviceName);
            throw new RuntimeException(e);
        }

    }
}
